import os
import time
from json import JSONDecodeError

import logging
from locust_plugins.connection_pools import FastHttpPool,RequestPool
from locust import HttpUser, constant, task, events, stats, run_single_user, FastHttpUser
import gevent
from locust.runners import STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP, MasterRunner, LocalRunner

def checker(environment):
    while not environment.runner.state in [STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP]:
        time.sleep(1)
        if environment.runner.stats.total.fail_ratio > 1:
            print(f"fail ratio was {environment.runner.stats.total.fail_ratio}, quitting")
            environment.runner.quit()
            return


@events.init.add_listener
def on_locust_init(environment, **_kwargs):
    # dont run this on workers, we only care about the aggregated numbers
    if isinstance(environment.runner, MasterRunner) or isinstance(environment.runner, LocalRunner):
        gevent.spawn(checker, environment)

@events.init_command_line_parser.add_listener
def _(parser):
    parser.add_argument("--my-argument", type=str, env_var="LOCUST_MY_ARGUMENT", default="", help="It's working")
    # Set `include_in_web_ui` to False if you want to hide from the web UI
    parser.add_argument("--my-ui-invisible-argument", include_in_web_ui=False, default="I am invisible")


@events.test_start.add_listener
def _(environment, **kw):
    print(f"Custom argument supplied: {environment.parsed_options.my_argument}")


class WebsiteUser(HttpUser):
    @task
    def my_task(self):
        print(f"my_argument={self.environment.parsed_options.my_argument}")
        print(f"my_ui_invisible_argument={self.environment.parsed_options.my_ui_invisible_argument}")


@events.request.add_listener
def my_request_handler(request_type, name, response_time, response_length, response,
                       context, exception, start_time, url, **kwargs):

    if exception:
        print(f"Request to {name} failed with exception {exception}")
    else:
        print(f"Successfully made a request to: {name}")
        print(f"The response was {response.text}")

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print("A new test is starting")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("A new test is ending")

@events.quitting.add_listener
def _(environment, **kw):
    if environment.stats.total.fail_ratio > 0.01:
        logging.error("Test failed due to failure ratio > 1%")
        environment.process_exit_code = 1
    elif environment.stats.total.avg_response_time > 200:
        logging.error("Test failed due to average response time ratio > 200 ms")
        environment.process_exit_code = 1
    elif environment.stats.total.get_response_time_percentile(0.95) > 800:
        logging.error("Test failed due to 95th percentile response time > 800 ms")
        environment.process_exit_code = 1
    else:
        environment.process_exit_code = 0

# @events.init.add_listener
# def on_locust_init(environment, **kw):
#     # You should now be able to start locust and browse to http://127.0.0.1:8089/added_page
#     @environment.web_ui.app.route("/added_page")
#     def my_added_page():
#         return "Another page"
# @events.test_start.add_listener
# def on_test_start(environment, **kwargs):
#     if not isinstance(environment.runner, MasterRunner):
#         print("Beginning test setup")
#     else:
#         print("Started test from Master node")
#
# @events.test_stop.add_listener
# def on_test_stop(environment, **kwargs):
#     if not isinstance(environment.runner, MasterRunner):
#         print("Cleaning up test data")
#     else:
#         print("Stopped test from Master node")





class RequestPoolExample(HttpUser):
    """HttpUser using requestPool"""

    def on_start(self):
        self.pool = RequestPool(user=self)

        # Alternatively, you can override the client attribute of the user
        # but if you have any type checking tools setup they will complain
        # self.client = RequestPool(user=self) # type: ignore

    @task
    def task_1(self):
        """example task"""
        self.pool.get(url="/")
        # self.client.get(url="/foo/bar")

    @task
    def task_2(self):
        """example task"""
        payload = {"bin": "baz"}
        self.pool.post(url="/foo/bar", data=payload)
        # self.client.get(url="/foo/bar", data=payload)


class FastHttpPoolExample(FastHttpUser):
    """FastHttpUser using"""

    def on_start(self):
        self.pool = FastHttpPool(user=self)
        # self.client = FastHttpPool(user=self) # type: ignore

    @task
    def task_1(self):
        """example task"""
        self.pool.get(path="/")
        # self.client.get(url="/foo/bar")

    @task
    def task_2(self):
        """example task"""
        payload = {"bin": "baz"}
        self.pool.post(path="/foo/bar", data=payload)
        # self.client.get(url="/foo/bar", data=payload)

class UploadFile(HttpUser):
    wait_time = constant(2)
    host = "https://www.retone-hearing.fit"
    header = {"Content-Type": "application/json"}
    #
    # fixed_count = 10

    # def on_start(self):
    #     self.pool = RequestPool(user=self)
    #
    # @task
    # def task_exap(self):
    #
    #     self.pool.post('/microServiceUser/needAppUserAuth/responseEncode/uploadUserProfile', files="xx",)



    @task(4)
    def upload_file(self):
        url = '/microServiceUser/needAppUserAuth/responseEncode/uploadUserProfile'

        with open(r"C:\Users\Administrator\Pictures\R-C.jpg", mode='rb') as image:
            image_rb = image.read()
        files = {'file': image_rb}
        with self.client.post(url, files=files, catch_response=True) as resp:
            try:
                if resp.status_code != 200 or resp.json()['data'] == None:
                    resp.failure(f"=======错误返回:{resp.text}==========")
                elif resp.elapsed.total_seconds() > 0.5:
                    resp.failure("Request took too long")
                else:
                    print(resp.json())
                    resp.success()
            except JSONDecodeError:
                resp.failure("Response could not be decoded as JSON")
            except KeyError:
                resp.failure("Response did not contain expected key 'greeting'")

    @task(1)
    def multiple_groupings_example(self):
        # Statistics for these requests will be grouped under: /blog/?id=[id]
        with self.client.rename_request("/blog?id=[id]"):
            for i in range(10):
                self.client.get("/blog?id=%i" % i)

        # Statistics for these requests will be grouped under: /article/?id=[id]
        with self.client.rename_request("/article?id=[id]"):
            for i in range(10):
                self.client.get("/article?id=%i" % i)


if __name__ == '__main__':
    os.system('locust -f locustfile_uploadfile.py --web-host="127.0.0.1"')
    # debug
    # run_single_user(UploadFile)

    